﻿import sqlite3
from dataclasses import dataclass
from typing import Optional


@dataclass
class Transition:
    episode_id: int
    state: tuple
    action: int
    next_state: tuple
    reward: float
    done: bool


class GameRecorder:
    def __init__(self, db_path="game.db"):
        self.db_path = db_path
        self.current_episode = self._get_max_episode() + 1
        self._init_db()

    def _init_db(self):
        """初始化数据库结构"""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS transitions (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    episode_id INTEGER NOT NULL,
                    state_x INTEGER NOT NULL,
                    state_y INTEGER NOT NULL,
                    action INTEGER NOT NULL,
                    next_state_x INTEGER NOT NULL,
                    next_state_y INTEGER NOT NULL,
                    reward REAL NOT NULL,
                    done BOOLEAN NOT NULL,
                    timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
                )
            """
            )
            conn.commit()

    def new_episode(self) -> int:
        """开始新的一局游戏"""
        self.current_episode = self._get_max_episode() + 1
        return self.current_episode

    def record(
        self, state: tuple, action: int, next_state: tuple, reward: float, done: bool
    ):
        """记录单次状态转移"""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute(
                """
                INSERT INTO transitions 
                (episode_id, state_x, state_y, action, next_state_x, next_state_y, reward, done)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            """,
                (
                    self.current_episode,
                    state[0],
                    state[1],
                    action,
                    next_state[0],
                    next_state[1],
                    reward,
                    done,
                ),
            )
            conn.commit()

    def _get_max_episode(self) -> int:
        """查询当前最大episode_id"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute("SELECT MAX(episode_id) FROM transitions")
            return cursor.fetchone()[0] or 0

    # 新增查询接口
    def get_episode_transitions(self, episode_id: int) -> list[Transition]:
        """获取指定episode的所有转移记录"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute(
                """
                SELECT state_x, state_y, action, next_state_x, next_state_y, reward, done
                FROM transitions
                WHERE episode_id = ?
                ORDER BY id
            """,
                (episode_id,),
            )

            return [
                Transition(
                    episode_id=episode_id,
                    state=(row[0], row[1]),
                    action=row[2],
                    next_state=(row[3], row[4]),
                    reward=row[5],
                    done=row[6],
                )
                for row in cursor.fetchall()
            ]

    def batch_record(self, transitions: list[Transition]):
        """批量记录状态转移"""
        with sqlite3.connect(self.db_path) as conn:
            conn.executemany(
                """
                INSERT INTO transitions 
                (episode_id, state_x, state_y, action, next_state_x, next_state_y, reward, done)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                """,
                [
                    (
                        t.episode_id,
                        t.state[0],
                        t.state[1],
                        t.action,
                        t.next_state[0],
                        t.next_state[1],
                        t.reward,
                        t.done,
                    )
                    for t in transitions
                ],
            )
            conn.commit()

    def get_transitions(self, limit=100, offset=0) -> list[Transition]:
        """分页获取所有转移记录"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute(
                """
                SELECT episode_id, state_x, state_y, action, 
                       next_state_x, next_state_y, reward, done
                FROM transitions
                ORDER BY id
                LIMIT ? OFFSET ?
                """,
                (limit, offset),
            )
            return [Transition(*row) for row in cursor.fetchall()]

    def get_all_episodes(self) -> list[int]:
        """获取所有episode ID列表"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute(
                "SELECT DISTINCT episode_id FROM transitions ORDER BY episode_id"
            )
            return [row[0] for row in cursor.fetchall()]


if __name__ == "__main__":
    # 查询记录
    recorder = GameRecorder()
    print("\n=== 游戏记录 ===")
    episodes = recorder.get_all_episodes()
    for ep in episodes:
        print(f"\n第 {ep} 局记录:")
        transitions = recorder.get_episode_transitions(ep)
        for t in transitions:
            print(
                f"状态: {t.state} → 动作: {t.action} → 新状态: {t.next_state} | "
                f"奖励: {t.reward} | 结束: {t.done}"
            )